//////////////////////////////////////////////////////////////////
// (c) copyright 2010- by jeongnim kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   National Center for Supercomputing Applications 
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//
// supported by 
//   national center for supercomputing applications, uiuc
//////////////////////////////////////////////////////////////////
// -*- c++ -*-
#ifndef QMCPLUSPLUS_DNS_H
#define QMCPLUSPLUS_DNS_H
#include <vector>
#include <complex>
#include <iostream>
#include <cstdlib>
#include <omp.h>

#define MEMBER_INLINE inline
//#define MEMBER_INLINE 

#define FUNCTION_INLINE inline
//#define FUNCTION_INLINE 

#if defined(HAVE_MKL)
#include <mkl_vsl.h>
#endif
#define USE_ALLOCATOR
#if defined(USE_ALLOCATOR)
inline void* allocate(size_t PS, size_t n)
{
  void *d;
  posix_memalign(&d,PS,n);
  return d;
}
#endif

/** definition of Timer */
struct Timer
{
  double start_time;
  inline Timer() { start_time=omp_get_wtime();}
  inline void restart() {start_time=omp_get_wtime();}
  inline double elapsed() const 
  {
    return omp_get_wtime()-start_time;
  }
};

namespace DNS
{
  template<typename T, typename RT>
    FUNCTION_INLINE void 
    axpy(T* restrict Q, const T* restrict rhs, RT one_third, int n)
    {
      for(int i=0; i<n; ++i) Q[i]+=one_third*rhs[i];
    }

  template<typename T, typename RT>
    FUNCTION_INLINE void 
    axpy(T* restrict Q, const T* restrict Q_old, const T* restrict rhs, RT delta, int n)
    {
      for(int i=0; i<n; ++i) Q[i] = Q_old[i] + delta*rhs[i];
    }

  /** aux function to evaluate vorticity_x in k-space for 3xyz storage
  */
  template<typename T, typename T2, typename ST>
    FUNCTION_INLINE void
    vorticity_x(T* restrict x, const T* restrict y, const T* restrict z
        , const T2* restrict n1_map, ST n1
        , const T2* restrict n2_map, ST n2
        , const T2* restrict n3_map, ST n3
        )
    {
      for(int i=0,ijk=0; i<n1; ++i)
      {
        for(int j=0; j<n2; ++j)
        {
          T jg=n2_map[j];
          //#pragma unroll(16)
          for(int k=0; k<n3; ++k, ijk+=2)
          {
            T kg=n3_map[k];
            x[ijk  ] = jg*z[ijk+1]-kg*y[ijk+1];
            x[ijk+1] = kg*y[ijk  ]-jg*z[ijk  ];
          }
        }
      }
    }

  /** aux function to evaluate vorticity_y in k-space for 3xyz storage
  */
  template<typename T, typename T2, typename ST>
    FUNCTION_INLINE void
    vorticity_y(T* restrict y, const T* restrict z, const T* restrict x
        , const T2* restrict n1_map, ST n1
        , const T2* restrict n2_map, ST n2
        , const T2* restrict n3_map, ST n3
        )
    {
      for(int i=0,ijk=0; i<n1; ++i)
      {
        T ig=n1_map[i];
        for(int j=0; j<n2; ++j)
        {
          //#pragma unroll(16)
          for(int k=0; k<n3; ++k, ijk+=2)
          {
            T kg=n3_map[k];
            y[ijk  ] = kg*x[ijk+1]-ig*z[ijk+1];
            y[ijk+1] = ig*z[ijk  ]-kg*x[ijk  ];
          }
        }
      }
    }

  /** aux function to evaluate vorticity_z in k-space for 3xyz storage
  */
  template<typename T, typename T2, typename ST>
    FUNCTION_INLINE void
    vorticity_z(T* restrict z, const T* restrict x, const T* restrict y
        , const T2* restrict n1_map, ST n1
        , const T2* restrict n2_map, ST n2
        , const T2* restrict n3_map, ST n3
        )
    {
      for(int i=0,ijk=0; i<n1; ++i)
      {
        T ig=n1_map[i];
        for(int j=0; j<n2; ++j)
        {
          T jg=n2_map[j];
          //#pragma unroll(16)
          for(int k=0; k<n3; ++k, ijk+=2)
          {
            z[ijk  ] = ig*y[ijk+1]-jg*x[ijk+1];
            z[ijk+1] = jg*x[ijk  ]-ig*y[ijk  ];
          }
        }
      }
    }

  template<typename T>
    FUNCTION_INLINE void 
    outerproduct3(const T* restrict rhsx, const T* restrict rhsy, const T* restrict rhsz
        ,T* restrict tx, T* restrict ty, T* restrict tz , size_t n)
    {
      for(size_t i=0; i<n;++i)
      {
        T vorx=rhsx[i]; 
        T vory=rhsy[i]; 
        T vorz=rhsz[i];
        T qx=tx[i]; 
        T qy=ty[i]; 
        T qz=tz[i];
        tx[i]=qy*vorz-qz*vory;
        ty[i]=qz*vorx-qx*vorz;
        tz[i]=qx*vory-qy*vorz;
      }
    }

  /** out += mkk * in
   * @param in input array 
   * @param out output array
   * @param mkk pre-computed laplacian factor
   * @param n size of mkk
   *
   * Note that in/out are complex data, i.e., in[2*n], out[2*n]
   */
  template<typename T>
    FUNCTION_INLINE void 
    multadd(const T* restrict in, const T* restrict mkk
        , T* restrict out, size_t n)
    {
      for(size_t i=0,i2=0; i<n;++i,i2+=2) 
      {
        out[i2] += mkk[i]*in[i2];
        out[i2+1] += mkk[i]*in[i2+1];
      }
    }

  template<typename T>
  struct Problem
  {
    /** @enum D=D */
    enum {D=3};
    ///global grid
    size_t dims_g[D];
    ///local grid
    size_t dims_l[D];
    ///first grid element
    size_t first[D];
    ///last grid element
    size_t last[D];
    ///number of doubles dims_l[0]*dims_l[1]*dims_l[2]
    size_t max_size;
    ///grid partition 
    int partitions[D];
    ///kmap[i] is the k_i for the i-th Cartesian coordinate
    std::vector<std::vector<T> > kmap;
    ///save -|k|^2
    std::vector<double> mkk;

    /** constructor for a cubic cell
     * @param n dimension 
     * @param npx distribution of the first dimension
     * @param npy distribution of the second dimension
     */
    inline Problem(size_t n, int npx, int npy)
    {
      for(int i=0;i<D; ++i) dims_g[i]=n;
      for(int i=0;i<D; ++i) partitions[i]=1;
      partitions[0]=npx;
      partitions[1]=npy;
      reset();
    }

    /** this can be anything: using 4 doubles */
    inline void reset()
    {
      for(int i=0; i<D; ++i) dims_l[i]=dims_g[i]/partitions[i];
      for(int i=0; i<D; ++i) dims_l[i]=dims_g[i]/partitions[i];
      size_t mypagesize=sysconf(_SC_PAGESIZE)/sizeof(T);
      max_size=dims_l[0]*dims_l[1]*dims_l[2];
      size_t n=(max_size/mypagesize)*mypagesize;
      max_size=(n<max_size)?n+mypagesize:n;
      if(kmap.empty()) kmap.resize(D);

      //this should be templated with the distribution method
      kmap[0].resize(dims_l[1]/2,0.11);
      kmap[1].resize(dims_l[0],0.13);
      kmap[2].resize(dims_l[2],0.15);
      mkk.resize(size()/2,-0.25);
    }

    inline size_t size() const 
    {
      return dims_l[0]*dims_l[1]*dims_l[2];
    }

    /** return the size of aligned data */
    inline size_t capacity() const
    {
      return max_size;
    }

    inline size_t size(int  i) const
    {
      return dims_l[i];
    }

    MEMBER_INLINE void vorticity(const T* restrict source, T* restrict target) const
    {
      vorticity_x(target,source+max_size,source+2*max_size
          , &(kmap[0][0]), kmap[0].size() , &(kmap[1][0]), kmap[1].size() , &(kmap[2][0]), kmap[2].size()
          );
      vorticity_y(target+max_size  ,source+2*max_size,source         
          , &(kmap[0][0]), kmap[0].size() , &(kmap[1][0]), kmap[1].size() , &(kmap[2][0]), kmap[2].size()
          );
      vorticity_z(target+max_size*2,source+max_size  ,source+max_size  
          , &(kmap[0][0]), kmap[0].size() , &(kmap[1][0]), kmap[1].size() , &(kmap[2][0]), kmap[2].size()
          );
    }

    /** target = source x target
    */
    MEMBER_INLINE void outerproduct(const T* restrict source, T* restrict target) const
    {
      outerproduct3(source,source+max_size,source+2*max_size
          , target, target+max_size, target+2*max_size
          , size());
    }

    /** target += -|k|^2 source
    */
    MEMBER_INLINE void laplacians(const T* restrict source, T* restrict target) const
    {
      multadd(source,           &mkk[0], target           , mkk.size());
      multadd(source+max_size  ,&mkk[0], target+max_size  , mkk.size());
      multadd(source+2*max_size,&mkk[0], target+2*max_size, mkk.size());
    }

    MEMBER_INLINE void aliasing(T* restrict rhs) 
    {
    }
  };

  template<typename T, unsigned NV=3>
    class Container
    {
      public:
        /** @typedef real data type */
        typedef T real_type;
        /** @typedef complex data type */
        typedef std::complex<T> complex_type;
        const Problem<T>& mygrid;
        T* mydata;
        size_t mysize;
        /** this will be determined by allocator */
        size_t myoffset;

        inline Container(const Problem<T>& in):mygrid(in)
      {
        //need smarter allocators
        myoffset=mygrid.capacity();
        mysize=myoffset*NV;
        //mysize=mygrid.size()*NV;
        size_t mypagesize=sysconf(_SC_PAGESIZE);
        mydata=static_cast<T*>(allocate(mypagesize,mysize*sizeof(T)));
      }

        ~Container()
        {
          if(mydata) free(mydata);
        }


        inline T* data() { return mydata;}
        inline const T* data() const { return mydata;}

        /** return the starting address of the variable i
        */
        inline T* data(int iv) { return mydata+myoffset*iv;}
        /** return the starting address of the variable i
        */
        inline const T* data(int iv) const { return mydata+myoffset*iv;}

        inline void randomize()
        {
#if defined(HAVE_MKL)
          VSLStreamStatePtr stream;  
          vslNewStream( &stream, VSL_BRNG_MCG31, 777 );
          vdRngUniform( VSL_METHOD_DUNIFORM_STD, stream, mysize, mydata, 0.0, 1.0 ); 
#else
          for(size_t i=0; i<mysize; ++i) mydata[i]=rand();
#endif
        }

        inline T* begin() { return mydata;}
        inline const T* begin() const { return mydata;}
        inline T* end() { return mydata+mysize;}
        inline const T* end() const { return mydata+mysize;}

        /** return the size with padding */
        inline size_t size() const { return mysize;}

        /** return the offset */
        inline size_t offset() const { return myoffset;}

        /** inplace inverse 3Dfft */
        inline void ifft() {}
        /** outplace inverse 3DFFT */
        inline void ifft(Container<T,NV>& a) const {}
        /** inplace forward 3Dfft */
        inline void fft() {}
        /** outplace forward 3Dfft */
        inline void fft(Container<T,NV>& a) const {}


      private:
        Container(const Container<T>& in):mygrid(in.mygrid),mydata(0)
        { }
    };
}
#endif

